import torch
from loss.base_loss import Loss

class MeanTeacherLoss(Loss):
    def __init__(self, name="mse", weight=2):
        super(MeanTeacherLoss, self).__init__()
        # Mean Teacher Loss is typically used in semi-supervised learning
        # It encourages the model to produce consistent predictions
        # for augmented versions of the same input.
        if name.lower() == "mse":
            self.loss_fn = torch.nn.MSELoss()
        elif name.lower() == "bce":
            self.loss_fn = torch.nn.BCELoss()
        else:
            raise ValueError(f"Unsupported loss type: {name}. Use 'mse' or 'bce'.")
        self.weight = weight
    
    def forward(self, student_preds: torch.Tensor, teacher_preds: torch.Tensor):
        """
        Compute the Mean Teacher Loss.
        
        Args:
            student_preds (torch.Tensor): Predictions from the student model.
            teacher_preds (torch.Tensor): Predictions from the teacher model.
        
        Returns:
            torch.Tensor: Computed loss value.
        """
        return self.weight * self.loss_fn(student_preds, teacher_preds)